{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hello World! NeMo-Run Style!\n",
    "\n",
    "Let's start with a simple notebook that demonstrates how to use NeMo-Run to configure and launch your Python functions. In this example, we take simple addition functions that look like:\n",
    "\n",
    "```python\n",
    "def add(a: int, b: int) -> int:\n",
    "    print(f\"Adding {a} to {b} returns {a + b}\")\n",
    "    return a + b\n",
    "```\n",
    "\n",
    "and use NeMo-Run to configure and launch it. This basic notebook demonstrates many of the building blocks in our NeMo-Run library.\n",
    "\n",
    "As described in the introduction, NeMo-Run is a tool that allows you to:\n",
    "1. Configure your functions or scripts in a Pythonic way.\n",
    "2. Launch them on any supported remote cluster directly from your local workstation.\n",
    "3. Manage them using `run.Experiment`.\n",
    "\n",
    "Let's get into it.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
    "# SPDX-License-Identifier: Apache-2.0\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "# Set up and imports\n",
    "import logging\n",
    "\n",
    "import nemo_run as run\n",
    "from simple.add import add\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s %(message)s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure the Python Function\n",
    "\n",
    "The first step in using NeMo-Run is to configure your Python function. As mentioned above, we're trying to configure the `add` function. Configuration is similar to Python's native [functools.partial](https://docs.python.org/3/library/functools.html#functools.partial). In fact, for most functions, you can replace `functools.partial` with `run.Partial` and it should still work. Configuration just ties your function and your arguments together to create a `run.Partial` object which can executed at a later time. This is done by building the `run.Partial` object which will recursively build any configured objects and then return a [`functools.partial` object](https://docs.python.org/3/library/functools.html#partial-objects).\n",
    "\n",
    "This brings us to the first building blocks of NeMo-Run: `Partial` and `Config`. These are buildables that allow you to configure functions, arguments, and objects in a Pythonic way. Under the hood, we use [fiddle](https://fiddle.readthedocs.io/en/latest/) to manage the configuration. Our `Partial` and `Config` classes are subclasses of `fiddle.Partial` and `fiddle.Config`, respectively, with additional features to enhance UX and programmability. We take some inspiration from [Praxis](https://github.com/google/praxis/blob/main/praxis/pax_fiddle.py#L72), but also have custom addons that should improve the user experience.\n",
    "\n",
    "We already discussed `run.Partial`. Similarly, `run.Config` takes a `fn` or `class` as the first argument, followed by the `fn`'s or `class`\\' `*args`, and `**kwargs` as the subsequent arguments. However, on building `run.Config`, it actually calls the underlying `fn` or, in case of a `class`, its `__init__` method with the arguments tied into the `Config`. For example:\n",
    "\n",
    "\n",
    "```python\n",
    "import fiddle as fdl\n",
    "import nemo_run as run\n",
    "def hello_world(msg: str):\n",
    "    print(f\"Hello World! {msg}\")\n",
    "\n",
    "cfg = run.Config(hello_world, msg=\"How are you?\")\n",
    "partial = run.Partial(hello_world, msg=\"How are you?\")\n",
    "\n",
    "fn = fdl.build(partial)\n",
    "fn()\n",
    "#>>> Hello World! How are you?\n",
    "\n",
    "built_cfg = fdl.build(cfg)\n",
    "#>>> Hello World! How are you?\n",
    "\n",
    "print(built_cfg is None)\n",
    "#>>> True\n",
    "```\n",
    "\n",
    "The `Partial` and `Config` classes also come with utilities that provide a visual representation of the object (if you have `graphviz` installed). Try running the cell below to configure the `add` function and visualize the configured `Partial`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = run.Partial(add, a=5, b=10)\n",
    "fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The configured function is now ready for execution. Before we proceed, let's look at configuring complex functions that take non-primitive types. Additionally, we will look at another utility that NeMo-Run provides: `run.autoconvert`. This decorator allows you to automatically convert functions that return regular Python objects into functions that return a `run.Config` of the underlying object. To demonstrate this, let's consider an `add_object` function that looks like:\n",
    "\n",
    "```python\n",
    "@dataclass\n",
    "class SomeObject:\n",
    "    value_1: int\n",
    "    value_2: int\n",
    "    value_3: int\n",
    "\n",
    "\n",
    "def add_object(obj_1: SomeObject, obj_2: SomeObject) -> SomeObject:\n",
    "    result = SomeObject(\n",
    "        value_1=obj_1.value_1 + obj_2.value_1,\n",
    "        value_2=obj_1.value_2 + obj_2.value_2,\n",
    "        value_3=obj_1.value_3 + obj_2.value_3,\n",
    "    )\n",
    "    print(f\"{result = }\")\n",
    "\n",
    "    return result\n",
    "```\n",
    "\n",
    "To configure this function, you can use `run.Partial`. Note, however, that you need to ensure the arguments to the function are `run.Config` or `run.Partial` instances. This is necessary for serialization and remote execution, but more on that later. For now, a basic configuration for add_object looks like:\n",
    "\n",
    "```python\n",
    "run.Partial(\n",
    "    add_object,\n",
    "    obj_1=run.Config(SomeObject, value_1=10, value_2=20, value_3=30),\n",
    "    obj_2=run.Config(SomeObject, value_1=10, value_2=20, value_3=30),\n",
    ")\n",
    "```\n",
    "\n",
    "Now, let's say you have a regular Python function that returns `SomeObject` and you want to use it as one of the arguments to `add_object`. You can decorate it using `run.autoconvert` as follows:\n",
    "```python\n",
    "@run.autoconvert\n",
    "def commonly_used_object() -> SomeObject:\n",
    "    return SomeObject(\n",
    "        value_1=5,\n",
    "        value_2=10,\n",
    "        value_3=15,\n",
    "    )\n",
    "```\n",
    "\n",
    "The `run.autoconvert` decorator uses fiddle's autoconfig and parses the AST so that you get the following:\n",
    "\n",
    "```python\n",
    "commonly_used_object() == run.Config(SomeObject, value_1=5, value_2=10, value_3=15)\n",
    "```\n",
    "\n",
    "\n",
    "You can then use the decorated function as follows:\n",
    "```python\n",
    "run.Partial(\n",
    "    add_object,\n",
    "    obj_1=run.Config(SomeObject, value_1=10, value_2=20, value_3=30),\n",
    "    obj_2=commonly_used_object(),\n",
    ")\n",
    "```\n",
    "\n",
    "Additionally, you can also use args in the function. Note, however, that `run.autoconvert` currently doesn't support control flow and may be unreliable for complex code. In that case, you can define a function to return the `run.Config` directly, as follows:\n",
    "\n",
    "```python\n",
    "def commonly_used_config() -> run.Config[SomeObject]:\n",
    "    config = run.Config(\n",
    "        SomeObject,\n",
    "        value_1=5,\n",
    "        value_2=10,\n",
    "        value_3=15,\n",
    "    )\n",
    "\n",
    "    for i in range(10):\n",
    "        config.value_1 *= i\n",
    "        config.value_2 += i\n",
    "        config.value_3 -= i\n",
    "\n",
    "    return config\n",
    "```\n",
    "\n",
    "> 📝 `autoconvert` returns a `run.Config` of the underlying object by default, whereas `@autoconvert(partial=True)` returns a `run.Partial` of the underlying object.\n",
    "\n",
    "One additional thing is that the arguments of a `run.Partial` or `run.Config` can be mutated via dot access, as shown below:\n",
    "```python\n",
    "fn.obj_1.value_1 = 100\n",
    "fn.obj_2.value_2 *= 2\n",
    "...\n",
    "```\n",
    "\n",
    "As you can see, our tool is designed to provide a lot of flexibility. Run the cells shown below to experiment with these APIs and learn about the different ways to configure your task or function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simple.add import add_object, commonly_used_object, commonly_used_object_2, SomeObject\n",
    "\n",
    "fn = run.Partial(\n",
    "    add_object,\n",
    "    obj_1=run.Config(SomeObject, value_1=10, value_2=20, value_3=30),\n",
    "    obj_2=run.Config(SomeObject, value_1=10, value_2=20, value_3=30),\n",
    ")\n",
    "fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = run.Partial(add_object, obj_1=commonly_used_object(), obj_2=commonly_used_object_2())\n",
    "fn.obj_1.value_1 = 100\n",
    "fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn = run.Partial(\n",
    "    add_object,\n",
    "    obj_1=commonly_used_object(),\n",
    "    obj_2=run.Config(SomeObject, value_1=10, value_2=20, value_3=30),\n",
    ")\n",
    "fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execute the Configured Function\n",
    "\n",
    "The previous section concludes the core features of the NeMo-Run configuration. Now, let's move on to execution. The core building block to execute a single function is `run.run`.\n",
    "\n",
    "The easiest way to execute a configured function is to execute it directly, just like you would execute a normal Python function. You can do this using the `direct` option in `run.run` as shown below:\n",
    "```python\n",
    "# You can also do a dryrun to see what's getting executed where.\n",
    "run.run(fn, direct=True, dryrun=True)\n",
    "run.run(fn, direct=True)\n",
    "```\n",
    "\n",
    "> 📝 As of now, the `run` function doesn't recreate the return value, so it is the user's responsibility to manage the artifacts from the function. For example, if you are launching a function to train an ML model, make sure you configure a job directory that you can access and inspect later. We are working on improving the management capabilities of a run and open to feedback, so please reach out to us if you have any ideas.\n",
    "\n",
    "Try it out in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can also do a dryrun to see what's getting executed where.\n",
    "run.run(fn, direct=True, dryrun=True)\n",
    "run.run(fn, direct=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Direct runs are not that interesting, as you wouldn't need to use `run` if you were only running the function directly. Beyond direct runs, we support a few executors out of the box:\n",
    "1. Local executor: This will run the function locally, but in a separate process. It is useful when you are using components like `torchrun` to launch your job/run.\n",
    "1. Slurm executor: This will run the function on a remote Slurm cluster. Currently, only Slurm clusters that have [Pyxis](https://github.com/NVIDIA/pyxis) are supported, but we plan to add support for other types of Slurm clusters in the future. Reach out to us if you have a specific request.\n",
    "1. Skypilot executor: This will run the function on any cloud supported by [Skypilot](https://skypilot.readthedocs.io/en/latest/).\n",
    "\n",
    "To learn more about executors, go to the execution guide linked in [README](../../README.md).\n",
    "\n",
    "> 📝 Currently, we do not pickle the function, but rely on Python module references for serializing the configured function. This means that executing a function defined in this notebook (or in the same script calling `run.run`) will probably not work. During packaging, we currently use `git archive` for our remote Slurm executor. You can specify a subpath which will be set as the working directory for the execution. To execute `add_object` function in this example, you need to provide a subpath of `examples/hello-world`, assuming `examples` is at the root of the repository. This will ensure that, during remote execution, imports of the style `from simple.add import add_object` work properly. We are looking to improve this user experience and open to suggestions, so please reach out to us if you have any ideas. We also plan to add `cloudpickle` support for arbitrarily defined functions.\n",
    "\n",
    "Under the hood, we are using [TorchX](https://pytorch.org/torchx/latest/) to manage the execution. However, this is fairly abstracted away from the user and we can potentially support more standalone executor libraries or add custom schedulers for TorchX in the future.\n",
    "\n",
    "Let's see the local executor in action. The local executor is the simplest executor we have and can be initialized without any arguments. Once we have an instance of `run.LocalExecutor`, we can pass it to `run.run` to execute the same configured function on the local executor, which will run the function in a separate process. This demonstrates the ease of use provided by NeMo-Run. You can configure a function once and execute it on any supported remote cluster seamlessly. Later, we will also explore `run.Experiment` which allows you to combine and manage multiple runs, providing additional flexibility to the user.\n",
    "\n",
    "Execute the cells below to run your function on the local executor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "executor = run.LocalExecutor()\n",
    "executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.run(fn, executor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you notice the logs, it has mentions of Experiment. Each `run.run` internally creates an experiment with a single task to provide management capabilities for a run out of the box. We also make the `run.Experiment` API available publicly to create custom experiments and workflows. Check out our next tutorial in this series named [`hello_experiments.ipynb`](hello_experiments.ipynb) to learn more about experiments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the same way, you can execute your configured function on a remote Slurm cluster or remote Skypilot cluster as well. \n",
    "\n",
    "Below, we show an example of defining a `run.SlurmExecutor`. NeMo-Run also sets up an `ssh` tunnel automatically for you to connect to the Slurm cluster and handle the packaging of the code so you never have to leave your local workstation. You can configure the Slurm executor as below:\n",
    "\n",
    "```python\n",
    "tunnel_cfg = run.Config(\n",
    "    TunnelConfig,\n",
    "    host=os.environ[\"SLURM_HOST\"],\n",
    "    user=os.environ[\"SLURM_USER\"],\n",
    "    job_dir=os.environ[\"SLURM_JOBDIR\"],\n",
    ")\n",
    "packager = run.Config(\n",
    "    GitArchivePackager,\n",
    "    use_torchrun=False,\n",
    "    subpath=\"examples/hello-world\"\n",
    ")\n",
    "\n",
    "executor = run.Config(\n",
    "    SlurmExecutorConfig,\n",
    "    account=os.environ[\"SLURM_ACCT\"],\n",
    "    partition=os.environ[\"SLURM_PARTITION\"],\n",
    "    nodes=1,\n",
    "    ntasks_per_node=1,\n",
    "    tunnel=tunnel,\n",
    "    container_image=os.environ[\"BASE_IMAGE\"],\n",
    "    time=\"00:30:00\",\n",
    "    packager=packager,\n",
    ")\n",
    "```\n",
    "\n",
    "More details about configuring the executors can be found in the `API Reference` portion of our docs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
